{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "cd262589",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "import torch\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import numpy as np\n",
    "from torch.utils.data.dataset import random_split"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "435699b0",
   "metadata": {},
   "source": [
    "https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "4b4a94ad",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.],\n",
       "        [1.],\n",
       "        [2.],\n",
       "        [3.],\n",
       "        [4.],\n",
       "        [5.],\n",
       "        [6.],\n",
       "        [7.],\n",
       "        [8.],\n",
       "        [9.]])"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.arange(10).float().unsqueeze(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "87f48a73",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0, 0, 0, 1, 1, 0, 1, 0, 1, 1])"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[torch.arange(1).float() for _ in range(10)]\n",
    "torch.randint(0, 2, (10,))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f74e229f",
   "metadata": {},
   "source": [
    "# 简单使用"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9b8d3e17",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(tensor([0.]), tensor(1))\n",
      "(tensor([1.]), tensor(1))\n",
      "(tensor([2.]), tensor(1))\n",
      "(tensor([3.]), tensor(0))\n",
      "(tensor([4.]), tensor(0))\n",
      "(tensor([5.]), tensor(1))\n",
      "(tensor([6.]), tensor(1))\n",
      "(tensor([7.]), tensor(1))\n",
      "(tensor([8.]), tensor(1))\n",
      "(tensor([9.]), tensor(0))\n"
     ]
    }
   ],
   "source": [
    "\n",
    "# 自定义数据集类\n",
    "class SimpleDataset(Dataset):\n",
    "    def __init__(self):\n",
    "        self.data = torch.arange(10).float().unsqueeze(1)\n",
    "        self.labels = torch.randint(0, 2, (10,))\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return self.data[idx], self.labels[idx]\n",
    "\n",
    "# 创建数据集对象\n",
    "dataset = SimpleDataset()\n",
    "for _ in dataset:\n",
    "    print(_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c540f8cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "dl = DataLoader(dataset,batch_size=3,shuffle=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "33980bb1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[tensor([[0.],\n",
      "        [1.],\n",
      "        [2.]]), tensor([1, 1, 1])]\n"
     ]
    }
   ],
   "source": [
    "for _ in dl:\n",
    "    print(_)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8e2a2c0",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "f3a7588b",
   "metadata": {},
   "source": [
    "# 简单使用2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "3dc9539e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(tensor([9, 5, 3, 2, 1]), tensor(0))\n",
      "(tensor([2, 4, 7, 2, 0]), tensor(1))\n",
      "(tensor([8, 2, 1, 8, 0]), tensor(1))\n",
      "(tensor([6, 5, 7, 3, 3]), tensor(0))\n",
      "(tensor([1, 1, 3, 9, 0]), tensor(0))\n",
      "(tensor([0, 5, 0, 8, 2]), tensor(0))\n",
      "(tensor([9, 9, 3, 5, 0]), tensor(0))\n",
      "(tensor([4, 1, 8, 5, 4]), tensor(1))\n",
      "(tensor([0, 5, 7, 9, 9]), tensor(0))\n",
      "(tensor([3, 9, 4, 2, 2]), tensor(1))\n"
     ]
    }
   ],
   "source": [
    "\n",
    "# 自定义数据集类\n",
    "class SimpleDataset(Dataset):\n",
    "    def __init__(self):\n",
    "        self.data = [torch.randint(0,10,size=(5,)) for _ in range(10)]\n",
    "        self.labels = torch.randint(0, 2, (10,))\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return self.data[idx], self.labels[idx]\n",
    "\n",
    "# 创建数据集对象\n",
    "dataset = SimpleDataset()\n",
    "for _ in dataset:\n",
    "    print(_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b48c1aae",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[tensor([[9, 5, 3, 2, 1],\n",
      "        [2, 4, 7, 2, 0],\n",
      "        [8, 2, 1, 8, 0]]), tensor([0, 1, 1])]\n"
     ]
    }
   ],
   "source": [
    "dl = DataLoader(dataset,batch_size=3,shuffle=False)\n",
    "for _ in dl:\n",
    "    print(_)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ee7b103",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "71707d22",
   "metadata": {},
   "source": [
    "# 简单使用3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "cd9d15d9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(tensor([[ 1.0787,  0.6038],\n",
      "        [-0.8529,  0.9280],\n",
      "        [-0.5963, -2.2425]]), tensor(1))\n",
      "(tensor([[-1.1323,  1.0617],\n",
      "        [-0.4451, -0.9041],\n",
      "        [-0.6046, -0.0604]]), tensor(1))\n",
      "(tensor([[-0.7629,  0.3164],\n",
      "        [-2.0978,  0.7600],\n",
      "        [-0.8037,  0.3958]]), tensor(0))\n",
      "(tensor([[ 1.3690,  2.1073],\n",
      "        [ 0.7149, -1.3104],\n",
      "        [ 0.2154, -0.7409]]), tensor(0))\n",
      "(tensor([[ 0.2746, -0.6653],\n",
      "        [ 0.4064,  0.4487],\n",
      "        [ 0.0369, -0.6497]]), tensor(1))\n",
      "(tensor([[-0.3606,  0.6560],\n",
      "        [-0.3733,  0.9193],\n",
      "        [ 0.5409, -2.0034]]), tensor(1))\n",
      "(tensor([[ 0.4379,  1.2955],\n",
      "        [ 0.3813, -0.2941],\n",
      "        [-0.2024,  0.4714]]), tensor(1))\n",
      "(tensor([[-1.4072,  1.7835],\n",
      "        [-0.2982, -1.5961],\n",
      "        [-1.0693,  0.6387]]), tensor(0))\n",
      "(tensor([[ 0.7840,  1.6168],\n",
      "        [-0.4520, -0.7149],\n",
      "        [ 0.6697, -2.2039]]), tensor(1))\n",
      "(tensor([[ 0.4413,  2.6311],\n",
      "        [ 0.5956,  1.2098],\n",
      "        [-1.4732,  0.2527]]), tensor(0))\n"
     ]
    }
   ],
   "source": [
    "# 自定义数据集类\n",
    "class SimpleDataset(Dataset):\n",
    "    def __init__(self):\n",
    "        self.data = [torch.randn(size=(3,2))  for _ in range(10)]\n",
    "        self.labels = torch.randint(0, 2, (10,))\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return self.data[idx], self.labels[idx]\n",
    "\n",
    "# 创建数据集对象\n",
    "dataset = SimpleDataset()\n",
    "for _ in dataset:\n",
    "    print(_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "47ea0e00",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[tensor([[[ 1.0787,  0.6038],\n",
      "         [-0.8529,  0.9280],\n",
      "         [-0.5963, -2.2425]],\n",
      "\n",
      "        [[-1.1323,  1.0617],\n",
      "         [-0.4451, -0.9041],\n",
      "         [-0.6046, -0.0604]],\n",
      "\n",
      "        [[-0.7629,  0.3164],\n",
      "         [-2.0978,  0.7600],\n",
      "         [-0.8037,  0.3958]]]), tensor([1, 1, 0])]\n"
     ]
    }
   ],
   "source": [
    "dl = DataLoader(dataset,batch_size=3,shuffle=False)\n",
    "for _ in dl:\n",
    "    print(_)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3639cad1",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "27a37214",
   "metadata": {},
   "source": [
    "# 简单使用3\n",
    "\n",
    "使用 [(),()]  list嵌套元祖的方式，也可以给Datalodaer\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "ca12ade5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([-0.21447456, -2.09040562,  0.26303774,  0.20256307, -0.5272343 ,\n",
       "       -1.160507  , -0.68202777, -0.36434035, -1.19781263, -0.76954088,\n",
       "       -0.63928313,  0.68477847])"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.random.randn(12,)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "8bd40c5e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(tensor([ 0.6813, -0.1519, -0.1112, -0.4410, -0.3830, -0.1806,  0.6537, -0.6138,\n",
       "          -0.1143, -0.8125,  0.2666, -0.1740]),\n",
       "  3),\n",
       " (tensor([-0.9116,  0.0507, -0.0310, -1.5456, -0.2881, -0.6581,  0.3059,  1.1918,\n",
       "           1.1074,  0.3890,  0.7147,  0.5896]),\n",
       "  3),\n",
       " (tensor([-1.6671,  0.9718, -0.5326,  1.6194,  1.2120,  1.4282, -0.7556, -0.0405,\n",
       "           0.9529, -0.0756,  1.7988, -0.7582]),\n",
       "  3),\n",
       " (tensor([-1.4463, -0.6748, -1.2180,  0.4149, -0.5986,  1.1878, -1.2538, -0.1086,\n",
       "           0.7726,  0.5414, -0.3686,  0.5352]),\n",
       "  3),\n",
       " (tensor([ 1.8083, -0.4097, -0.2901,  2.0756, -1.0595, -0.6754,  1.1753,  0.1619,\n",
       "           0.2109, -2.1750,  0.9192,  0.5339]),\n",
       "  3),\n",
       " (tensor([-0.7375, -0.2307, -0.1567,  0.4722, -0.2873, -1.2887, -0.7509,  2.1771,\n",
       "           1.0626,  0.8531, -0.8413, -0.3236]),\n",
       "  3),\n",
       " (tensor([-1.5112,  1.4696,  0.1573, -2.0331,  1.1367, -0.7673,  0.7631,  0.3858,\n",
       "          -0.6006, -0.6368, -0.6449, -0.0601]),\n",
       "  3),\n",
       " (tensor([-0.8224,  0.0333, -0.4454,  0.7060, -1.6265,  0.0325, -0.1390, -1.1727,\n",
       "          -0.0303,  0.3982,  0.5769, -1.1676]),\n",
       "  3),\n",
       " (tensor([ 0.2477,  0.4333,  0.1312,  1.0274,  0.1921,  2.3538,  0.7166, -0.2415,\n",
       "          -0.4750, -1.1216, -0.1775, -1.0767]),\n",
       "  3),\n",
       " (tensor([ 0.1046,  0.7464,  0.1151,  0.3229,  0.4511, -1.0861,  1.0270,  0.1710,\n",
       "          -1.2483, -0.4887, -0.6558, -0.2149]),\n",
       "  3)]"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset = [(torch.randn(size=(12,)),3) for _ in range(10)]\n",
    "# dataset = [(np.random.randn(12,),3) for _ in range(10)]\n",
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9a69e16",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(tensor([-0.5655, -0.6495,  1.4332, -2.3340,  1.0154,  1.7060, -0.7024, -0.1839,\n",
       "           0.1985, -0.3578,  0.8015,  0.5620]),\n",
       "  3),\n",
       " (tensor([-0.1807, -0.9340,  0.8804, -0.1583, -0.3176, -0.0210, -0.5074, -0.0806,\n",
       "           0.5647, -0.7087,  0.5722, -0.5843]),\n",
       "  3),\n",
       " (tensor([-1.7092,  1.3454,  1.5122,  0.8428,  0.0266,  0.9123, -0.0587,  0.3546,\n",
       "          -1.8058, -0.5775, -1.7202,  0.9026]),\n",
       "  3),\n",
       " (tensor([-0.5502, -0.4609,  0.6004,  0.9399, -0.8653, -1.9673, -0.0234,  0.8523,\n",
       "           0.6454, -0.1925,  0.3064,  0.5776]),\n",
       "  3),\n",
       " (tensor([-0.7640, -0.4797,  0.7678,  0.6625,  0.9420, -0.6987,  0.7405, -0.4923,\n",
       "           0.3404, -0.4353,  1.8158, -0.6178]),\n",
       "  3),\n",
       " (tensor([-0.1634, -0.7753, -0.6942, -0.6523, -0.1644,  0.2456, -0.6178,  1.1081,\n",
       "          -1.0165,  1.2135,  0.3222,  0.1127]),\n",
       "  3),\n",
       " (tensor([ 2.1152, -0.5901, -0.4098,  1.0767, -3.1121, -0.5476,  1.2146,  1.6605,\n",
       "           0.9829, -2.2070, -0.5402, -1.1961]),\n",
       "  3),\n",
       " (tensor([-1.3582,  0.8612,  0.7711, -1.5103,  1.7021, -0.5656, -0.0400, -0.7425,\n",
       "          -0.2605, -1.1486, -0.7739, -0.0678]),\n",
       "  3),\n",
       " (tensor([ 0.0868,  0.0450,  0.7537, -0.2641,  0.2864, -0.0252, -1.0505, -0.6640,\n",
       "           1.1312,  1.3226, -0.6320, -0.8879]),\n",
       "  3)]"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "train_len = int(len(dataset) * 0.95)\n",
    "sub_train_, sub_valid_ = random_split(\n",
    "    dataset, [train_len, len(dataset) - train_len])\n",
    "list(sub_train_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "854803a6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[tensor([[-0.5655, -0.6495,  1.4332, -2.3340,  1.0154,  1.7060, -0.7024, -0.1839,\n",
      "          0.1985, -0.3578,  0.8015,  0.5620],\n",
      "        [-0.1807, -0.9340,  0.8804, -0.1583, -0.3176, -0.0210, -0.5074, -0.0806,\n",
      "          0.5647, -0.7087,  0.5722, -0.5843],\n",
      "        [-1.7092,  1.3454,  1.5122,  0.8428,  0.0266,  0.9123, -0.0587,  0.3546,\n",
      "         -1.8058, -0.5775, -1.7202,  0.9026]]), tensor([3, 3, 3])]\n",
      "[tensor([[-0.5502, -0.4609,  0.6004,  0.9399, -0.8653, -1.9673, -0.0234,  0.8523,\n",
      "          0.6454, -0.1925,  0.3064,  0.5776],\n",
      "        [-0.7640, -0.4797,  0.7678,  0.6625,  0.9420, -0.6987,  0.7405, -0.4923,\n",
      "          0.3404, -0.4353,  1.8158, -0.6178],\n",
      "        [-0.1634, -0.7753, -0.6942, -0.6523, -0.1644,  0.2456, -0.6178,  1.1081,\n",
      "         -1.0165,  1.2135,  0.3222,  0.1127]]), tensor([3, 3, 3])]\n",
      "[tensor([[ 2.1152, -0.5901, -0.4098,  1.0767, -3.1121, -0.5476,  1.2146,  1.6605,\n",
      "          0.9829, -2.2070, -0.5402, -1.1961],\n",
      "        [-1.3582,  0.8612,  0.7711, -1.5103,  1.7021, -0.5656, -0.0400, -0.7425,\n",
      "         -0.2605, -1.1486, -0.7739, -0.0678],\n",
      "        [ 0.0868,  0.0450,  0.7537, -0.2641,  0.2864, -0.0252, -1.0505, -0.6640,\n",
      "          1.1312,  1.3226, -0.6320, -0.8879]]), tensor([3, 3, 3])]\n"
     ]
    }
   ],
   "source": [
    "dl = DataLoader(sub_train_,batch_size=3,shuffle=False)\n",
    "for _ in dl:\n",
    "    print(_)\n",
    "    # break"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cad8431",
   "metadata": {},
   "source": [
    "collate_fn 自定义处理函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "51659c5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def change(batchs):\n",
    "    '''\n",
    "    batch 就是\n",
    "    [(tensor([ 0.9818,  0.1940,  0.0436, -0.8141,  1.4636]), 3), (tensor([-0.0840, -0.0065, -2.6674, -0.3946,  2.8424]), 3), (tensor([ 2.2204, -1.0966, -0.2777,  0.3998,  1.1074]), 3)]\n",
    "    '''\n",
    "    # print(batchs)\n",
    "    data=[]\n",
    "    label = []\n",
    "    for batch in batchs:\n",
    "        data.extend(batch[0])\n",
    "        label.append(batch[1])\n",
    "    # print('data',data)\n",
    "    return  torch.tensor(data),torch.tensor(label)\n",
    "        \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4bd57944",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(tensor([ 0.6813, -0.1519, -0.1112, -0.4410, -0.3830, -0.1806,  0.6537, -0.6138,\n",
      "        -0.1143, -0.8125,  0.2666, -0.1740, -0.9116,  0.0507, -0.0310, -1.5456,\n",
      "        -0.2881, -0.6581,  0.3059,  1.1918,  1.1074,  0.3890,  0.7147,  0.5896,\n",
      "        -1.6671,  0.9718, -0.5326,  1.6194,  1.2120,  1.4282, -0.7556, -0.0405,\n",
      "         0.9529, -0.0756,  1.7988, -0.7582]), tensor([3, 3, 3]))\n"
     ]
    },
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mThe Kernel crashed while executing code in the current cell or a previous cell. \n",
      "\u001b[1;31mPlease review the code in the cell(s) to identify a possible cause of the failure. \n",
      "\u001b[1;31mClick <a href='https://aka.ms/vscodeJupyterKernelCrash'>here</a> for more info. \n",
      "\u001b[1;31mView Jupyter <a href='command:jupyter.viewOutput'>log</a> for further details."
     ]
    }
   ],
   "source": [
    "dl = DataLoader(dataset,batch_size=3,shuffle=False,collate_fn=change)\n",
    "for _ in dl:\n",
    "    print(_)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffd9618b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "472a457a",
   "metadata": {},
   "source": [
    "# 演示4   IterableDataset\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b8f860bf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([6, 9, 2, 0, 1])\n",
      "tensor([4, 9, 9, 6, 1])\n",
      "tensor([1, 0, 9, 7, 4])\n",
      "tensor([9, 0, 4, 0, 1])\n",
      "tensor([5, 9, 9, 3, 7])\n",
      "tensor([0, 7, 5, 9, 7])\n",
      "tensor([2, 1, 8, 1, 2])\n",
      "tensor([8, 0, 8, 3, 9])\n",
      "tensor([4, 6, 6, 9, 9])\n",
      "tensor([6, 1, 2, 7, 3])\n",
      "tensor([[3, 2, 4, 0, 8],\n",
      "        [0, 0, 0, 1, 4],\n",
      "        [0, 5, 8, 0, 4],\n",
      "        [0, 2, 0, 3, 2]])\n"
     ]
    }
   ],
   "source": [
    "import math,torch\n",
    "from torch.utils.data import IterableDataset,DataLoader,TensorDataset\n",
    "class MyIterableDataset(torch.utils.data.IterableDataset):\n",
    "    def __init__(self, start, end):\n",
    "        super(MyIterableDataset).__init__()\n",
    "        assert end > start, \"this example code only works with end >= start\"\n",
    "        self.start = start\n",
    "        self.end = end\n",
    "    def __iter__(self):\n",
    "        worker_info = torch.utils.data.get_worker_info()\n",
    "        if worker_info is None:  # single-process data loading, return the full iterator\n",
    "            iter_start = self.start\n",
    "            iter_end = self.end\n",
    "        else:  # in a worker process\n",
    "            # split workload\n",
    "            per_worker = int(math.ceil((self.end - self.start) / float(worker_info.num_workers)))\n",
    "            worker_id = worker_info.id\n",
    "            iter_start = self.start + worker_id * per_worker\n",
    "            iter_end = min(iter_start + per_worker, self.end)\n",
    "        s = torch.randint(0,10,(10,5))\n",
    "        # return iter((s,s))\n",
    "        # return iter(range(iter_start, iter_end))\n",
    "        return iter(torch.randint(0,10,(10,5)))\n",
    "# should give same set of data as range(3, 7), i.e., [3, 4, 5, 6].\n",
    "ds = MyIterableDataset(start=3, end=7)\n",
    "for y in ds:\n",
    "    print(y)\n",
    "# Single-process loading\n",
    "s=DataLoader(ds, num_workers=0,batch_size=4)\n",
    "for _ in s:\n",
    "    print(_)\n",
    "    break\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ef65e18",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torch_py38",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
